package dk.au.cs.peterbn.processing.compilationJobs;

import dk.au.cs.peterbn.patterns.TypeRole;
import dk.au.cs.peterbn.processing.ProcessingJob;
import dk.au.cs.peterbn.processing.ProcessingUtil;

import javax.annotation.processing.RoundEnvironment;
import javax.lang.model.element.*;
import javax.lang.model.type.TypeMirror;
import javax.tools.Diagnostic;
import java.util.*;

import static javax.lang.model.util.ElementFilter.*;

/**
 * Assert that all subclasses of "abstract" pattern type roles are annotated with
 * the proper pattern annotation. This check should really help in ensuring that
 * all members of a pattern gets annotated.
 */
public class CheckAnnotatedSubClasses implements ProcessingJob {

    private final RoundEnvironment env;
    private final ProcessingUtil util;

    /**
     * Create a new CheckAnnotatedSubClasses ProcessingJob.
     *
     * @param env The current RoundEnvironment
     * @param util The ProcessingUtil
     */
    public CheckAnnotatedSubClasses(RoundEnvironment env, ProcessingUtil util) {
        this.env = env;
        this.util = util;
    }

    @Override
    public void process() {
        Set<? extends Element> rootElements = env.getRootElements();
        Set<TypeElement> typeElements =  getAllCompilationTypes(rootElements);
        for (TypeElement compilationClass : typeElements) {
            List<AnnotationMirror> patternAnnotations = util.getPatternAnnotationMirrors(compilationClass);
            Set<ParentAnnotation> parentPatternAnnotations = getParentPatternAnnotations(compilationClass);
            List<TypeRole> unImplemented = new ArrayList<TypeRole>();
            for (ParentAnnotation parentPatternAnnotation : parentPatternAnnotations) {
                TypeElement parentElement = parentPatternAnnotation.getAnnotee();
                AnnotationMirror parentAnnotation = parentPatternAnnotation.getAnnotation();
                if (!compilationClass.equals(parentElement)) {
                    TypeRole typeRole = util.getTypeRole(parentElement, parentAnnotation);
                    Set<TypeRole> implementers = typeRole.getImplementers();
                    boolean implementsOneRole = false;
                    for (AnnotationMirror patternAnnotation : patternAnnotations) {
                        TypeRole compilationClassRole = util.getTypeRole(compilationClass, patternAnnotation);
                        implementsOneRole |= implementers.remove(compilationClassRole);
                    }
                    if (!implementsOneRole) {
                        //Only report errors if we aren't implementing any roles
                        unImplemented.addAll(implementers);
                    }
                }
            }
            for (Iterator<TypeRole> iterator = unImplemented.iterator(); iterator.hasNext(); ) {
                TypeRole typeRole = iterator.next();
                boolean correctTarget = isCorrectTarget(compilationClass, typeRole);
                if (!correctTarget) {
                    iterator.remove();
                }

            }
            if (!unImplemented.isEmpty()) {
                for (TypeRole typeRole : unImplemented) {
                    util.printMessage(Diagnostic.Kind.ERROR, "Missing TypeRole: " + typeRole.getName(), compilationClass);
                }
            }
        }
    }

    private boolean isCorrectTarget(TypeElement compilationClass, TypeRole typeRole) {
        ElementKind kind = compilationClass.getKind();
        boolean isInterface = kind.isInterface();
        boolean isAbstractClass = kind.isClass() && compilationClass.getModifiers().contains(Modifier.ABSTRACT);
        boolean isConcreteClass = kind.isClass() && !compilationClass.getModifiers().contains(Modifier.ABSTRACT);
        return  TypeRole.Target.ANY == typeRole.getTarget() ||
                (typeRole.getTarget() == TypeRole.Target.ABSTRACT && (isInterface || isAbstractClass))
                || (typeRole.getTarget() == TypeRole.Target.CONCRETE && isConcreteClass);
    }

    private Set<ParentAnnotation> getParentPatternAnnotations(TypeElement compilationClass) {
        Set<ParentAnnotation> parentTypeRoles = new HashSet<ParentAnnotation>();
        if (compilationClass == null) {
            return parentTypeRoles;
        }

        List<AnnotationMirror> patternAnnotations = util.getPatternAnnotationMirrors(compilationClass);
        for (AnnotationMirror patternAnnotation : patternAnnotations) {
            parentTypeRoles.add(new ParentAnnotation(compilationClass, patternAnnotation));
        }

        TypeMirror superclass = compilationClass.getSuperclass();
        if (superclass != null) {
            TypeElement typeElement = util.asTypeElement(superclass);
            parentTypeRoles.addAll(getParentPatternAnnotations(typeElement));
        }
        List<? extends TypeMirror> interfaces = compilationClass.getInterfaces();
        for (TypeMirror anInterface : interfaces) {
            TypeElement typeElement = util.asTypeElement(anInterface);
            parentTypeRoles.addAll(getParentPatternAnnotations(typeElement));
        }
        return parentTypeRoles;
    }


    /**
     * Get all types in the current compilation unit.
     * We cannot simply get them from the set of PatternInstances, as it only includes types that are already
     * annotated with Pattern. We want all the types
     * @param rootElements The elements to search in
     * @return The set of type elements in the current compilation
     */
    private Set<TypeElement> getAllCompilationTypes(Collection<? extends Element> rootElements) {
        Set<TypeElement> typeElements = new HashSet<TypeElement>();
        List<TypeElement> typeElementList = typesIn(rootElements);
        for (TypeElement element : typeElementList) {
            typeElements.add(element);
            List<? extends Element> enclosedElements = element.getEnclosedElements();
            typeElements.addAll(getAllCompilationTypes(enclosedElements));
        }
        return typeElements;
    }

    private static class ParentAnnotation {
        final private TypeElement annotee;
        final private AnnotationMirror annotation;

        public ParentAnnotation(TypeElement annotee, AnnotationMirror annotation) {
            this.annotee = annotee;
            this.annotation = annotation;
        }

        public TypeElement getAnnotee() {
            return annotee;
        }

        public AnnotationMirror getAnnotation() {
            return annotation;
        }
    }
}
